#include "include/OpticalFlowProcessor.h"

OpticalFlowProcessor::OpticalFlowProcessor(int width, int height)
{
	displacement = Point2f(0,0);
	orientation = 0;
	prevFrame = new Mat(width, height, CV_8UC1);
	currFrame = new Mat(width, height, CV_8UC1);
	frameSize = Point(width,height);
	frameCentroid = Point2f(width / 2.0f, height /2.0f);
	winSize = Size(OPTICAL_FLOW_WINSIZE, OPTICAL_FLOW_WINSIZE);
	srand(time(NULL));
}

OpticalFlowProcessor::~OpticalFlowProcessor()
{
	prevFrame->release();
	currFrame->release();
}

void OpticalFlowProcessor::Process(int action, long long newFrameAddr, long long addrFrameToDraw)
{

	Mat* newFrame = (Mat*)newFrameAddr;
	Mat* outFrame = (Mat*)addrFrameToDraw;

	if (action == OPTICAL_FLOW_ACTION_CALC_SIFT_FEATURES)
	{
		/*
		SIFT sift;
		vector<KeyPoint> keypoints;

		detector->detect(*currFrame,keypoints);
		Ptr<DescriptorExtractor> featureExtractor = DescriptorExtractor::create("SIFT");

		Mat descriptors;
		featureExtractor->compute(*currFrame, keypoints, descriptors);

		/*
		Mat outputFrame;
		drawKeypoints(*outFrame, keypoints, outputFrame, Scalar(255,255,0), DrawMatchesFlags::DEFAULT);
		outputFrame.copyTo(*outFrame);
		*/
	}
	else if (action == OPTICAL_FLOW_ACTION_PLAY_AROUND)
	{

		OrbFeatureDetector detector(OPTICAL_FLOW_MAX_FEATURES);
		FREAK freak;

		newFrame->copyTo(*currFrame);
		detector.detect(*currFrame, newKeyPoints);
		freak.compute(*currFrame, newKeyPoints, newDescriptors);
		if (newDescriptors.type() != CV_32F)
		{
			newDescriptors.convertTo(newDescriptors,CV_32F);
		}



		if (oldKeyPoints.size() > 0 )
		{
			/*
			FlannBasedMatcher matcher;
			vector<DMatch> matches;
			matcher.match(newDescriptors, oldDescriptors, matches);
			Mat tempImg;
			drawMatches(*prevFrame, oldKeyPoints, *currFrame, newKeyPoints, matches, tempImg);
			transpose(tempImg, tempImg);
			resize(tempImg, *outFrame, Size(320,240));
			*/


		}

		currFrame->copyTo(*prevFrame);
		oldKeyPoints = newKeyPoints;
		newDescriptors.copyTo(oldDescriptors);
	}
	else
	{

		currFrame->copyTo(*prevFrame);
		newFrame->copyTo(*currFrame);

		//Detect Features using FAST features
		FastFeatureDetector detector(OPTICAL_FLOW_MAX_FEATURES);
		vector<KeyPoint> v;
		detector.detect(*newFrame, v);

		//calculate optical flow
		lastFeatures = currFeatures;
		currFeatures.clear();
		for (int i = 0; i < v.size(); i++)currFeatures.push_back(((KeyPoint)v[i]).pt);

		if (action == OPTICAL_FLOW_ACTION_FEATURES_ONLY)
		{
			for (int i = 0; i < currFeatures.size(); i++)
				circle(*outFrame, ((KeyPoint)(v[i])).pt, 10, Scalar(255,255,255,255));
		}
		else if (lastFeatures.size() && currFeatures.size())
		{
			vector<uchar> status;
			vector<float> err;

			calcOpticalFlowPyrLK(*prevFrame,*currFrame,lastFeatures,currFeatures,status,err,winSize);
			int validFeatureCount = 0;
			for (int i = 0; i < currFeatures.size(); i++)
			{
				if (status[i])
				{
					validFeatureCount++;
				}
			}

			if (validFeatureCount < OPTICAL_FLOW_RESET_THRESHOLD)
			{
				lastFeatures = currFeatures;
				currFrame->copyTo(*prevFrame);
			}
			else if (action == OPTICAL_FLOW_ACTION_CALC_FLOW_ONLY)
			{
				for (int i = 0; i < currFeatures.size(); i++)
				{
					DrawArrow(outFrame, ((Point2f)(lastFeatures[i])),((Point2f)(currFeatures[i])),1);
				}
			}
			else if (validFeatureCount > OPTICAL_FLOW_CALC_RANSAC_THRESHOLD)
			{

				//Perform Ransac to find out orientation and displacement;

				vector<Point2f> angleDistance;
				vector<Point2f> angleDistanceNorm;
				Point2f angleDistanceSum = Point2f(0,0);
				Point2f angleDistanceSumSq = Point2f(0,0);
				float maxInliner;
				int maxIndex;

				for (int i = 0;i < currFeatures.size(); i++)
				{
					float dx = ((Point2f)lastFeatures[i]).x - ((Point2f)currFeatures[i]).x;
					float dy = ((Point2f)lastFeatures[i]).y - ((Point2f)currFeatures[i]).y;
					float angle = atan2(dy,dx);
					float distance = sqrt (dy*dy + dx * dx);

					angleDistance.push_back(Point2f(angle,distance));
					angleDistanceSum.x += angle;
					angleDistanceSum.y += distance;
					angleDistanceSumSq.x += angle * angle;
					angleDistanceSumSq.y += distance * distance;

				}

				Point2f angleDistanceMean, angleDistanceVariance;
				angleDistanceMean.x = angleDistanceSum.x / (float)angleDistance.size();
				angleDistanceMean.y = angleDistanceSum.y / (float)angleDistance.size();
				angleDistanceVariance.x = angleDistanceSumSq.x / (float)angleDistance.size();
				angleDistanceVariance.y = angleDistanceSumSq.y / (float)angleDistance.size();
				angleDistanceVariance -= angleDistanceMean;

				for (int i = 0; i < currFeatures.size(); i++)
				{
					Point2f norm;
					norm.x = ((Point2f)(angleDistance[i])).x;
					norm.y = (((Point2f)(angleDistance[i])).y - angleDistanceMean.y) / angleDistanceVariance.y;
					angleDistanceNorm.push_back(norm);
				}

				for (int tryCount = 0; tryCount < OPTICAL_FLOW_MAX_RANSAC_TRY_COUNT; tryCount++)
				{
					int randIndex = rand() % angleDistanceNorm.size();
					Point2f testPoint = angleDistanceNorm[randIndex];
					int inlinerCount = 0;
					for (int i = 0;i < angleDistance.size(); i++)
					{
						float dx = testPoint.x - ((Point2f)angleDistanceNorm[i]).x;
						float dy = testPoint.x - ((Point2f)angleDistanceNorm[i]).x;

						if (abs(dx) > CV_PI)
						{
							dx = dx - 2 * CV_PI;
						}

						float error = dx * dx + dy * dy;
						if (error < OPTICAL_FLOW_MAX_RANSAC_ERROR)
						{
							inlinerCount++;
						}
					}

					if (inlinerCount > maxInliner)
					{
						maxInliner = inlinerCount;
						maxIndex = randIndex;
					}
				}

				displacement.x = ((Point2f)lastFeatures[maxIndex]).x - ((Point2f)currFeatures[maxIndex]).x;
				displacement.y = ((Point2f)lastFeatures[maxIndex]).y - ((Point2f)currFeatures[maxIndex]).y;
				orientation = ((Point2f)angleDistance[maxIndex]).x;

				if (action == OPTICAL_FLOW_ACTION_CALC_FLOW_RANSAC)
				{
					DrawArrow(outFrame,frameCentroid,frameCentroid + displacement, 2);
				}
			}
		}
	}
}

double OpticalFlowProcessor::GetDisplacementX()
{
	return displacement.x;
}

double OpticalFlowProcessor::GetDisplacementY()
{
	return displacement.y;
}

double OpticalFlowProcessor::GetOrientation()
{
	return orientation;
}

OpticalFlowProcessor* OpticalFlowProcessor::ofp = NULL;

OpticalFlowProcessor* OpticalFlowProcessor::InitProcessor(int width, int height)
{
	if (ofp != NULL)
	{
		Release();
	}

	ofp = new OpticalFlowProcessor(width, height);
	return ofp;
}

void OpticalFlowProcessor::DrawArrow(Mat* image, Point2f a, Point2f b, int thickness)
{
	Point p,q;
	p.x = (int)(a.x);
	p.y = (int)(a.y);
	q.x = (int)(b.x);
	q.y = (int)(b.y);

	double angle = atan2((double)p.y-q.y, (double)p.x-q.x);
	double hypotenuse = sqrt((p.y-q.y)* (p.y-q.y) + (p.x-q.x) * (p.x-q.x));
	q.x = (int) (p.x - 3 * hypotenuse * cos(angle));
	q.y = (int) (p.y - 3 * hypotenuse * sin(angle));
	line(*image, p, q, Scalar(255,255,255), thickness, CV_AA, 0);

	p.x = (int) (q.x + 9 * cos(angle + PI / 4));
	p.y = (int) (q.y + 9 * sin(angle + PI / 4));
	line(*image, p, q, Scalar(255,255,255), thickness, CV_AA, 0);

	p.x = (int) (q.x + 9 * cos(angle - PI / 4));
	p.y = (int) (q.y + 9 * sin(angle - PI / 4));
	line(*image, p, q, Scalar(255,255,255), thickness, CV_AA, 0);
}

bool OpticalFlowProcessor::IsInitialized()
{
	return ofp != NULL;
}

void OpticalFlowProcessor::Release()
{
	if (ofp != NULL)
	{
		delete ofp;
		ofp = NULL;
	}
}

